from pathlib import Path
from os import remove
from functools import partial
from shutil import move
from shutil import copyfile
from itertools import groupby
from itertools import islice
from concurrent.futures import wait
from video_compressor import GROUP_NUM
from video_compressor import THREAD_POOL
from video_compressor import printf
from video_compressor import Colors
from video_compressor import VIDEO_WIDE
from video_compressor import VIDEO_HIGH
from video_compressor import IN_DIR_PATH
from video_compressor import OUT_DIR_PATH
from video_compressor import VIDEO_EXTS
from video_compressor import get_file
from video_compressor.utils import get_video_info
from video_compressor.ffmpeg_video import ffmpy_gpu_video_splicing
from video_compressor.ffmpeg_video import run


def create_temp(future, dir):
    if future or future.result():
        for file in filter(lambda f: f.is_file() and f.suffix in VIDEO_EXTS and 'temp' in f.name, dir.iterdir()):
            remove(file)


def get_tasks():
    tasks, file_set = [], set()
    file_iter = get_file(IN_DIR_PATH)
    for file in file_iter:
        out_dir_path = OUT_DIR_PATH.joinpath(file.relative_to(IN_DIR_PATH))
        if not out_dir_path.parent.is_dir(): out_dir_path.parent.mkdir(parents=True, exist_ok=True)
        if out_dir_path.suffix not in VIDEO_EXTS:
            copyfile(file, out_dir_path)
        else: file_set.add(file)

    test = sorted(file_set, key=lambda x: x, reverse=False)
    for name, file_iter in groupby(test, lambda f: f.parent.stem):
        out_dir_path = OUT_DIR_PATH.joinpath(name)
        out_path = out_dir_path / f'{name}.mp4'
        if out_path.is_file(): continue
        base_list = list(file_iter)
        wide, high, total_time = get_video_info(base_list[0])
        video_wide = VIDEO_WIDE if VIDEO_WIDE else wide
        video_high = VIDEO_HIGH if VIDEO_HIGH else high
        if len(base_list) < GROUP_NUM:
            ffmpeg = splicing_video(name, base_list, out_dir_path, video_wide, video_high)
        else:
            ffmpeg = splicing_video_temp(name, base_list, out_dir_path, video_wide, video_high, GROUP_NUM)
        task = THREAD_POOL.submit(run, ffmpeg, out_path)
        task.add_done_callback(partial(create_temp, dir=out_dir_path))
        tasks.append(task)
    return tasks


def splicing_video_temp(name, base_list, out_dir_path: Path, video_wide, video_high, group_num=5):
    temp_tasks = []
    base_dir = out_dir_path / 'temp'
    if not base_dir.is_dir(): base_dir.mkdir(parents=True, exist_ok=True)
    lt = iter(base_list)
    groups = iter(lambda: list(islice(lt, group_num)), [])
    for i, group in enumerate(groups, start=1):
        out_path_ = base_dir / f'{i:03d}_temp.mp4'
        printf(f'{Colors.OKGREEN}创建任务{Colors.ENDC}{Colors.OKBLUE}{out_path_}{Colors.ENDC} {Colors.OKGREEN}递归拼接视频{Colors.ENDC}')
        ffmpeg_ = ffmpy_gpu_video_splicing(group, out_path_, video_wide, video_high)
        temp_tasks.append(THREAD_POOL.submit(run, ffmpeg_, out_path_))
    if temp_tasks: wait(temp_tasks)
    create_temp(True, out_dir_path)
    file_list = []
    for file in base_dir.iterdir():
        move(file, out_dir_path)
        file_list.append(out_dir_path / file.name)
    test = sorted(file_list, key=lambda x: x, reverse=False)
    if len(test) >= group_num:
        ffmpeg = splicing_video_temp(name, test, out_dir_path, video_wide, video_high, group_num)
    else:
        ffmpeg = splicing_video(name, test, out_dir_path, video_wide, video_high)
    return ffmpeg


def splicing_video(name, base_list, out_dir_path: Path, video_wide, video_high):
    out_path = out_dir_path / f'{name}.mp4'
    printf(f'{Colors.OKGREEN}创建任务{Colors.ENDC}{Colors.OKBLUE}{out_path}{Colors.ENDC} {Colors.OKGREEN}递归拼接视频{Colors.ENDC}')
    ffmpeg = ffmpy_gpu_video_splicing(base_list, out_path, video_wide, video_high)
    return ffmpeg

